library(RItools)
library(Matching)
library(coin)
library(exactRankTests)
library(foreign)
source("code/balanceCheck.R")

load("data/SMRPreplication.RData")  ## pending Veterans' Administration approval

vte <- c("gender", "phq1mdddx", "pssi1_tot", "age", "phq1tot", "fas1tot", "stroopINTtscore1", "lm1_t1recallSCALED", "lm2_t1recallSCALED")
vn <- c("Gender (1=male)", "Probable Depression", "PTSD Severity", "Age", "Depression Severity", "Verbal Fluency", "Executive Control", "Log Mem - Immediate", "Log Mem - Delayed")

# 9 was removed before 10 was randomized
# 19 was in for 20-24, out before 25
# 34 was in for 35-38, out before 39
# 45 was in for 46-51
# 47 was in for 48-51
### Therefore:
xc <- x[!(x$id %in% c(9, 19, 34, 45, 47)), ]

## Figure 6:
outc <- balance.check(dat = xc, vars.to.eval = vte, vars.type = c(rep("disc", 2), rep("cont", 7)), vars.names = vn, cex.pval.axis = .6, tr.var = "condition", plot.file.qq = "balbothc.pdf", plot.file.pvals = "balpvaluesc.pdf", xy.line.color = "grey", d2p.color = "grey")

## Prepare for dynamic balance
## calc p-value with exactly units used when current unit assigned
ps <- array(NA)
# make formula:
xv <- NULL
for(xxxv in 1:(length(vte)-1)){
  xv <- paste(xv, vte[xxxv], "+", sep="")
}
xv <- paste(xv, vte[length(vte)], sep="")
ffff <- as.formula(paste("I(as.numeric(as.factor(condition))) ~ ", xv, sep=""))

## Calculate dynamic balance
ps3 <- array(NA)
for(i in 3:nrow(x)){
	dat <- x[1:i, ]
	if(10 %in% dat$id){dat <- dat[!(dat$id %in% 9), ]}
	if(25 %in% dat$id){dat <- dat[!(dat$id %in% 19), ]}
	if(39 %in% dat$id){dat <- dat[!(dat$id %in% 34), ]}
	ps3[i] <- xBalance(ffff, data = dat, report = c("chisquare.test"))$overall$p.value
}
rm(vn, vte, xxxv)


## Figure 5:
den <- 40
ccc <- .45
hhh <- .05
plot(-10, -10, xlim = c(0, 51), ylim = c(0,1), xlab = "", ylab = "", xaxt = "n")
axis(1, at = c(1,10,20,30,40,50,51), lwd=0, lwd.ticks =1)
abline(v=9, col = "grey")
rect(19, 0, 24, 1, col = "grey", density = den)
rect(34, 0, 38, 1, col = "grey", density = den)
rect(45, 0, 51, 1, col = "grey", density = den)
rect(47, 0, 51, 1, col = "grey", density = den, angle = 315)
par(new = TRUE)
plot(ps3, xlim = c(0,51), ylim = c(0,1), type = "b", axes=FALSE, ylab = bquote(Blocked ~ d^2 ~ "p-values"), xlab = "Unit Sequence", ann = FALSE)
mtext(side = 2, text = bquote(Blocked ~ d^2 ~ "p-values"), line = 2.5)
mtext(side = 1, text = "Unit Sequence", line = 2.5)
text(11.5, hhh, "    9 enters, 
     leaves  ", cex =ccc)
arrows(10.2, hhh, 9.2, hhh, length = .05)
text(21.5, hhh, "19 included", cex = ccc)
text(36, hhh, "34 included", cex = ccc)
text(48, hhh+.05, "45 included", cex = ccc)
text(49, hhh, "47 included", cex = ccc)

rm(den, ccc, hhh)

## SB versus CR in PTSD:
store.crp <- NULL
xc.cr <- xc
set.seed(125)
for(i in 1:10000){
  xc.cr$condition <- sample(c("anagrams", "memory"), nrow(xc.cr), replace = TRUE)
  store.crp[i] <- xBalance(ffff, data = xc.cr, report = c("chisquare.test"))$overall$p.value
}
mean(store.crp < outc$d2.p)

## Supplementary Figure 4:
plot(density(outc$d2.p - store.crp), main="", xlab = "(SB p-value) - (10,000 CR p-values)")
abline(v=0, col = "grey")
text(.5, .2, "SB more balanced than 99% of CRs")

rm(xc, xc.cr, ps, ps3, xv, ffff, i, dat, outc, store.crp)
